﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace NanoMvvm.Collections
{
    public class SortedKeyedList<TKey, TVal> : SortedList<TVal>
    {
        private readonly Func<TVal, TKey> keySelector;

        public SortedKeyedList(Func<TVal, TKey> keySelector)
            : base(new DelegateComparer<TVal>((v1, v2) => Comparer<TKey>.Default.Compare(keySelector(v1), keySelector(v2))))
        {
            this.keySelector = keySelector;
        }

        public IEnumerable<TVal> SelectByKey(TKey key)
        {
            var result = SelectInterval(key, key);

            var count1 = result.Count();
            var count2 = ListInternal.Count(i => Equals(keySelector(i), key));

            if (count1 != count2)
            {
                Debugger.Break();
            }

            return SelectInterval(key, key);
        }

        public IEnumerable<TVal> SelectInterval(TKey low, TKey high, IntervalOptions options = IntervalOptions.Inclusive)
        {
            if (Comparer<TKey>.Default.Compare(low, high) > 0)
            {
                throw new IndexOutOfRangeException();
            }

            var index = BinarySearch(low);

            if (index < 0)
            {
                index = ~index;
            }

            if (options == IntervalOptions.Exclusive || options == IntervalOptions.StartExclusiveEndInclusive)
            {
                while (index < Count && Comparer<TKey>.Default.Compare(keySelector(ListInternal[index]), low) == 0)
                {
                    index++;
                }
            }

            while (index < Count && Comparer<TKey>.Default.Compare(keySelector(ListInternal[index]), high) < 0)
            {
                yield return ListInternal[index];
                index++;
            }

            if (options == IntervalOptions.Inclusive || options == IntervalOptions.StartExclusiveEndInclusive)
            {
                while (index < Count && Comparer<TKey>.Default.Compare(keySelector(ListInternal[index]), high) == 0)
                {
                    yield return ListInternal[index];
                    index++;
                }
            }
        }

        public IEnumerable<TVal> SelectGreater(TKey key)
        {
            var index = BinarySearch(key);

            if (index < 0)
            {
                yield break;
            }

            while (index < Count && Comparer<TKey>.Default.Compare(keySelector(ListInternal[index]), key) == 0)
            {
                index++;
            }

            while (index < Count)
            {
                yield return ListInternal[index];
                index++;
            }
        }

        public IEnumerable<TVal> SelectLesser(TKey key)
        {
            var index = BinarySearch(key);

            if (index < 0)
            {
                yield break;
            }

            index--;

            while (index >= 0)
            {
                yield return ListInternal[index];
                index--;
            }
        }

        public IEnumerable<TVal> SelectGreaterOrEqual(TKey key)
        {
            var index = BinarySearch(key);

            if (index < 0)
            {
                yield break;
            }

            while (index < Count)
            {
                yield return ListInternal[index];
                index++;
            }
        }

        public IEnumerable<TVal> SelectLesserOrEqual(TKey key)
        {
            var index = BinarySearch(key);

            if (index < 0)
            {
                yield break;
            }

            while (index + 1 < Count && Comparer<TKey>.Default.Compare(keySelector(ListInternal[index + 1]), key) == 0)
            {
                index++;
            }

            while (index >= 0)
            {
                yield return ListInternal[index];
                index--;
            }
        }

        private int BinarySearch(TKey key)
        {
            int min = 0;
            int max = Count - 1;

            while (min <= max)
            {
                int mid = (min + max) / 2;
                int comparison = Comparer<TKey>.Default.Compare(keySelector(ListInternal[mid]), key);

                if (comparison == 0)
                {
                    while (mid - 1 >= 0 && Comparer<TKey>.Default.Compare(keySelector(ListInternal[mid - 1]), key) == 0)
                    {
                        mid--;
                    }

                    return mid;
                }

                if (comparison < 0)
                {
                    min = mid + 1;
                }
                else
                {
                    max = mid - 1;
                }
            }

            return ~min;
        }
    }
}